{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2024 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dfsDR_omdNea"
      },
      "source": [
        "# Gemma - Integrate with OneTwo\n",
        "\n",
        "This notebook demonstrates how you can use Gemma through [OneTwo](https://github.com/google-deepmind/onetwo). OneTwo is a Python library designed to simplify interactions with LLMs and is released by Google DeepMind.\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/gemma-cookbook/blob/main/Gemma/[Gemma_1]Using_with_OneTwo.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FaqZItBdeokU"
      },
      "source": [
        "## Setup\n",
        "\n",
        "### Select the Colab runtime\n",
        "To complete this tutorial, you'll need to have a Colab runtime with sufficient resources to run the Gemma model. In this case, you can use a T4 GPU:\n",
        "\n",
        "1. In the upper-right of the Colab window, select **▾ (Additional connection options)**.\n",
        "2. Select **Change runtime type**.\n",
        "3. Under **Hardware accelerator**, select **T4 GPU**.\n",
        "\n",
        "### Gemma setup\n",
        "\n",
        "To complete this tutorial, you'll first need to complete the setup instructions at [Gemma setup](https://ai.google.dev/gemma/docs/setup). The Gemma setup instructions show you how to do the following:\n",
        "\n",
        "* Get access to Gemma on [Kaggle](kaggle.com).\n",
        "* Generate and configure a [Kaggle username and API key](https://www.kaggle.com/docs/api). Set them up as Colab secrets 'KAGGLE_USERNAME' and 'KAGGLE_KEY'."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A9sUQ4WrP-Yr"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from google.colab import userdata\n",
        "\n",
        "# Note: `userdata.get` is a Colab API. If you're not using Colab, set the env\n",
        "# vars as appropriate for your system.\n",
        "os.environ[\"KAGGLE_USERNAME\"] = userdata.get(\"KAGGLE_USERNAME\")\n",
        "os.environ[\"KAGGLE_KEY\"] = userdata.get(\"KAGGLE_KEY\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZjSga1qnNdGh"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: kaggle in /usr/local/lib/python3.10/dist-packages (1.6.14)\n",
            "Requirement already satisfied: six>=1.10 in /usr/local/lib/python3.10/dist-packages (from kaggle) (1.16.0)\n",
            "Requirement already satisfied: certifi>=2023.7.22 in /usr/local/lib/python3.10/dist-packages (from kaggle) (2024.2.2)\n",
            "Requirement already satisfied: python-dateutil in /usr/local/lib/python3.10/dist-packages (from kaggle) (2.8.2)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from kaggle) (2.31.0)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from kaggle) (4.66.4)\n",
            "Requirement already satisfied: python-slugify in /usr/local/lib/python3.10/dist-packages (from kaggle) (8.0.4)\n",
            "Requirement already satisfied: urllib3 in /usr/local/lib/python3.10/dist-packages (from kaggle) (2.0.7)\n",
            "Requirement already satisfied: bleach in /usr/local/lib/python3.10/dist-packages (from kaggle) (6.1.0)\n",
            "Requirement already satisfied: webencodings in /usr/local/lib/python3.10/dist-packages (from bleach->kaggle) (0.5.1)\n",
            "Requirement already satisfied: text-unidecode>=1.3 in /usr/local/lib/python3.10/dist-packages (from python-slugify->kaggle) (1.3)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->kaggle) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->kaggle) (3.7)\n",
            "Requirement already satisfied: kagglehub in /usr/local/lib/python3.10/dist-packages (0.2.5)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from kagglehub) (24.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from kagglehub) (2.31.0)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from kagglehub) (4.66.4)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->kagglehub) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->kagglehub) (3.7)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->kagglehub) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->kagglehub) (2024.2.2)\n"
          ]
        }
      ],
      "source": [
        "# Install kaggle\n",
        "! pip install kaggle\n",
        "! pip install kagglehub"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HdTPjx_sHPUG"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Kaggle credentials set.\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "5bc98192c4d04446950daf04d8323995",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "VBox(children=(HTML(value='<center> <img\\nsrc=https://www.kaggle.com/static/images/site-logo.png\\nalt=\\'Kaggle…"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import kagglehub\n",
        "from kagglehub import auth\n",
        "\n",
        "auth.set_kaggle_credentials(\n",
        "    username=userdata.get(\"KAGGLE_USERNAME\"), api_key=userdata.get(\"KAGGLE_KEY\")\n",
        ")\n",
        "kagglehub.login()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "06e_SNhjI0Jd"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "VARIANT = \"2b\"  # @param ['2b', '1.1-2b-it', '7b', '1.1-7b-it'] {type:\"string\"}\n",
        "weights_dir = kagglehub.model_download(f\"google/gemma/Flax/{VARIANT}\")\n",
        "\n",
        "checkpoint_path = os.path.join(weights_dir, VARIANT)\n",
        "vocab_path = os.path.join(weights_dir, \"tokenizer.model\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dn3BirzvHglS"
      },
      "source": [
        "## Use Gemma with OneTwo\n",
        "\n",
        "### Install OneTwo"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "77GlsVYpHoFN"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting git+https://github.com/google-deepmind/onetwo\n",
            "  Cloning https://github.com/google-deepmind/onetwo to /tmp/pip-req-build-zjyeelh4\n",
            "  Running command git clone --filter=blob:none --quiet https://github.com/google-deepmind/onetwo /tmp/pip-req-build-zjyeelh4\n",
            "  Resolved https://github.com/google-deepmind/onetwo to commit 43c3b67e4b479e49e31baefb304d1487b4c893bb\n",
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting gemma@ git+https://github.com/google-deepmind/gemma.git (from onetwo==0.1.0)\n",
            "  Cloning https://github.com/google-deepmind/gemma.git to /tmp/pip-install-8hrdcipz/gemma_c35bc1f9a46943db8dce9b89f9a3baed\n",
            "  Running command git clone --filter=blob:none --quiet https://github.com/google-deepmind/gemma.git /tmp/pip-install-8hrdcipz/gemma_c35bc1f9a46943db8dce9b89f9a3baed\n",
            "  Resolved https://github.com/google-deepmind/gemma.git to commit a24194737dcb54b7392091e9ba772aea1cb68ffb\n",
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: absl-py in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (2.1.0)\n",
            "Requirement already satisfied: aenum in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (3.1.15)\n",
            "Requirement already satisfied: dataclasses-json in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (0.6.6)\n",
            "Requirement already satisfied: fastapi in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (0.111.0)\n",
            "Requirement already satisfied: google-cloud-aiplatform in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (1.52.0)\n",
            "Requirement already satisfied: google-generativeai in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (0.5.4)\n",
            "Requirement already satisfied: html5lib in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (1.1)\n",
            "Requirement already satisfied: immutabledict in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (4.2.0)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (3.1.4)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (1.25.2)\n",
            "Requirement already satisfied: openai in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (1.30.2)\n",
            "Requirement already satisfied: pillow in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (9.4.0)\n",
            "Requirement already satisfied: pytest in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (7.4.4)\n",
            "Requirement already satisfied: portpicker in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (1.5.2)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (6.0.1)\n",
            "Requirement already satisfied: termcolor in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (2.4.0)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (4.66.4)\n",
            "Requirement already satisfied: typing_extensions in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (4.11.0)\n",
            "Requirement already satisfied: uvicorn in /usr/local/lib/python3.10/dist-packages (from onetwo==0.1.0) (0.29.0)\n",
            "Requirement already satisfied: marshmallow<4.0.0,>=3.18.0 in /usr/local/lib/python3.10/dist-packages (from dataclasses-json->onetwo==0.1.0) (3.21.2)\n",
            "Requirement already satisfied: typing-inspect<1,>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from dataclasses-json->onetwo==0.1.0) (0.9.0)\n",
            "Requirement already satisfied: starlette<0.38.0,>=0.37.2 in /usr/local/lib/python3.10/dist-packages (from fastapi->onetwo==0.1.0) (0.37.2)\n",
            "Requirement already satisfied: pydantic!=1.8,!=1.8.1,!=2.0.0,!=2.0.1,!=2.1.0,<3.0.0,>=1.7.4 in /usr/local/lib/python3.10/dist-packages (from fastapi->onetwo==0.1.0) (2.7.1)\n",
            "Requirement already satisfied: fastapi-cli>=0.0.2 in /usr/local/lib/python3.10/dist-packages (from fastapi->onetwo==0.1.0) (0.0.4)\n",
            "Requirement already satisfied: httpx>=0.23.0 in /usr/local/lib/python3.10/dist-packages (from fastapi->onetwo==0.1.0) (0.27.0)\n",
            "Requirement already satisfied: python-multipart>=0.0.7 in /usr/local/lib/python3.10/dist-packages (from fastapi->onetwo==0.1.0) (0.0.9)\n",
            "Requirement already satisfied: ujson!=4.0.2,!=4.1.0,!=4.2.0,!=4.3.0,!=5.0.0,!=5.1.0,>=4.0.1 in /usr/local/lib/python3.10/dist-packages (from fastapi->onetwo==0.1.0) (5.10.0)\n",
            "Requirement already satisfied: orjson>=3.2.1 in /usr/local/lib/python3.10/dist-packages (from fastapi->onetwo==0.1.0) (3.10.3)\n",
            "Requirement already satisfied: email_validator>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from fastapi->onetwo==0.1.0) (2.1.1)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->onetwo==0.1.0) (2.1.5)\n",
            "Requirement already satisfied: flax>=0.8 in /usr/local/lib/python3.10/dist-packages (from gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (0.8.3)\n",
            "Requirement already satisfied: sentencepiece<0.2.0,>=0.1.99 in /usr/local/lib/python3.10/dist-packages (from gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (0.1.99)\n",
            "Requirement already satisfied: google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0dev,>=1.34.1 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform->onetwo==0.1.0) (2.11.1)\n",
            "Requirement already satisfied: google-auth<3.0.0dev,>=2.14.1 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform->onetwo==0.1.0) (2.27.0)\n",
            "Requirement already satisfied: proto-plus<2.0.0dev,>=1.22.0 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform->onetwo==0.1.0) (1.23.0)\n",
            "Requirement already satisfied: protobuf!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.19.5 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform->onetwo==0.1.0) (3.20.3)\n",
            "Requirement already satisfied: packaging>=14.3 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform->onetwo==0.1.0) (24.0)\n",
            "Requirement already satisfied: google-cloud-storage<3.0.0dev,>=1.32.0 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform->onetwo==0.1.0) (2.8.0)\n",
            "Requirement already satisfied: google-cloud-bigquery!=3.20.0,<4.0.0dev,>=1.15.0 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform->onetwo==0.1.0) (3.21.0)\n",
            "Requirement already satisfied: google-cloud-resource-manager<3.0.0dev,>=1.3.3 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform->onetwo==0.1.0) (1.12.3)\n",
            "Requirement already satisfied: shapely<3.0.0dev in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform->onetwo==0.1.0) (2.0.4)\n",
            "Requirement already satisfied: docstring-parser<1 in /usr/local/lib/python3.10/dist-packages (from google-cloud-aiplatform->onetwo==0.1.0) (0.16)\n",
            "Requirement already satisfied: google-ai-generativelanguage==0.6.4 in /usr/local/lib/python3.10/dist-packages (from google-generativeai->onetwo==0.1.0) (0.6.4)\n",
            "Requirement already satisfied: google-api-python-client in /usr/local/lib/python3.10/dist-packages (from google-generativeai->onetwo==0.1.0) (2.84.0)\n",
            "Requirement already satisfied: six>=1.9 in /usr/local/lib/python3.10/dist-packages (from html5lib->onetwo==0.1.0) (1.16.0)\n",
            "Requirement already satisfied: webencodings in /usr/local/lib/python3.10/dist-packages (from html5lib->onetwo==0.1.0) (0.5.1)\n",
            "Requirement already satisfied: anyio<5,>=3.5.0 in /usr/local/lib/python3.10/dist-packages (from openai->onetwo==0.1.0) (3.7.1)\n",
            "Requirement already satisfied: distro<2,>=1.7.0 in /usr/lib/python3/dist-packages (from openai->onetwo==0.1.0) (1.7.0)\n",
            "Requirement already satisfied: sniffio in /usr/local/lib/python3.10/dist-packages (from openai->onetwo==0.1.0) (1.3.1)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from portpicker->onetwo==0.1.0) (5.9.5)\n",
            "Requirement already satisfied: iniconfig in /usr/local/lib/python3.10/dist-packages (from pytest->onetwo==0.1.0) (2.0.0)\n",
            "Requirement already satisfied: pluggy<2.0,>=0.12 in /usr/local/lib/python3.10/dist-packages (from pytest->onetwo==0.1.0) (1.5.0)\n",
            "Requirement already satisfied: exceptiongroup>=1.0.0rc8 in /usr/local/lib/python3.10/dist-packages (from pytest->onetwo==0.1.0) (1.2.1)\n",
            "Requirement already satisfied: tomli>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from pytest->onetwo==0.1.0) (2.0.1)\n",
            "Requirement already satisfied: click>=7.0 in /usr/local/lib/python3.10/dist-packages (from uvicorn->onetwo==0.1.0) (8.1.7)\n",
            "Requirement already satisfied: h11>=0.8 in /usr/local/lib/python3.10/dist-packages (from uvicorn->onetwo==0.1.0) (0.14.0)\n",
            "Requirement already satisfied: idna>=2.8 in /usr/local/lib/python3.10/dist-packages (from anyio<5,>=3.5.0->openai->onetwo==0.1.0) (3.7)\n",
            "Requirement already satisfied: dnspython>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from email_validator>=2.0.0->fastapi->onetwo==0.1.0) (2.6.1)\n",
            "Requirement already satisfied: typer>=0.12.3 in /usr/local/lib/python3.10/dist-packages (from fastapi-cli>=0.0.2->fastapi->onetwo==0.1.0) (0.12.3)\n",
            "Requirement already satisfied: jax>=0.4.19 in /usr/local/lib/python3.10/dist-packages (from flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (0.4.26)\n",
            "Requirement already satisfied: msgpack in /usr/local/lib/python3.10/dist-packages (from flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (1.0.8)\n",
            "Requirement already satisfied: optax in /usr/local/lib/python3.10/dist-packages (from flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (0.2.2)\n",
            "Requirement already satisfied: orbax-checkpoint in /usr/local/lib/python3.10/dist-packages (from flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (0.4.4)\n",
            "Requirement already satisfied: tensorstore in /usr/local/lib/python3.10/dist-packages (from flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (0.1.45)\n",
            "Requirement already satisfied: rich>=11.1 in /usr/local/lib/python3.10/dist-packages (from flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (13.7.1)\n",
            "Requirement already satisfied: googleapis-common-protos<2.0.dev0,>=1.56.2 in /usr/local/lib/python3.10/dist-packages (from google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0dev,>=1.34.1->google-cloud-aiplatform->onetwo==0.1.0) (1.63.0)\n",
            "Requirement already satisfied: requests<3.0.0.dev0,>=2.18.0 in /usr/local/lib/python3.10/dist-packages (from google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0dev,>=1.34.1->google-cloud-aiplatform->onetwo==0.1.0) (2.31.0)\n",
            "Requirement already satisfied: grpcio<2.0dev,>=1.33.2 in /usr/local/lib/python3.10/dist-packages (from google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0dev,>=1.34.1->google-cloud-aiplatform->onetwo==0.1.0) (1.64.0)\n",
            "Requirement already satisfied: grpcio-status<2.0.dev0,>=1.33.2 in /usr/local/lib/python3.10/dist-packages (from google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0dev,>=1.34.1->google-cloud-aiplatform->onetwo==0.1.0) (1.48.2)\n",
            "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from google-auth<3.0.0dev,>=2.14.1->google-cloud-aiplatform->onetwo==0.1.0) (5.3.3)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.10/dist-packages (from google-auth<3.0.0dev,>=2.14.1->google-cloud-aiplatform->onetwo==0.1.0) (0.4.0)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.10/dist-packages (from google-auth<3.0.0dev,>=2.14.1->google-cloud-aiplatform->onetwo==0.1.0) (4.9)\n",
            "Requirement already satisfied: google-cloud-core<3.0.0dev,>=1.6.0 in /usr/local/lib/python3.10/dist-packages (from google-cloud-bigquery!=3.20.0,<4.0.0dev,>=1.15.0->google-cloud-aiplatform->onetwo==0.1.0) (2.3.3)\n",
            "Requirement already satisfied: google-resumable-media<3.0dev,>=0.6.0 in /usr/local/lib/python3.10/dist-packages (from google-cloud-bigquery!=3.20.0,<4.0.0dev,>=1.15.0->google-cloud-aiplatform->onetwo==0.1.0) (2.7.0)\n",
            "Requirement already satisfied: python-dateutil<3.0dev,>=2.7.2 in /usr/local/lib/python3.10/dist-packages (from google-cloud-bigquery!=3.20.0,<4.0.0dev,>=1.15.0->google-cloud-aiplatform->onetwo==0.1.0) (2.8.2)\n",
            "Requirement already satisfied: grpc-google-iam-v1<1.0.0dev,>=0.12.4 in /usr/local/lib/python3.10/dist-packages (from google-cloud-resource-manager<3.0.0dev,>=1.3.3->google-cloud-aiplatform->onetwo==0.1.0) (0.13.0)\n",
            "Requirement already satisfied: certifi in /usr/local/lib/python3.10/dist-packages (from httpx>=0.23.0->fastapi->onetwo==0.1.0) (2024.2.2)\n",
            "Requirement already satisfied: httpcore==1.* in /usr/local/lib/python3.10/dist-packages (from httpx>=0.23.0->fastapi->onetwo==0.1.0) (1.0.5)\n",
            "Requirement already satisfied: annotated-types>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from pydantic!=1.8,!=1.8.1,!=2.0.0,!=2.0.1,!=2.1.0,<3.0.0,>=1.7.4->fastapi->onetwo==0.1.0) (0.7.0)\n",
            "Requirement already satisfied: pydantic-core==2.18.2 in /usr/local/lib/python3.10/dist-packages (from pydantic!=1.8,!=1.8.1,!=2.0.0,!=2.0.1,!=2.1.0,<3.0.0,>=1.7.4->fastapi->onetwo==0.1.0) (2.18.2)\n",
            "Requirement already satisfied: mypy-extensions>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from typing-inspect<1,>=0.4.0->dataclasses-json->onetwo==0.1.0) (1.0.0)\n",
            "Requirement already satisfied: httptools>=0.5.0 in /usr/local/lib/python3.10/dist-packages (from uvicorn->onetwo==0.1.0) (0.6.1)\n",
            "Requirement already satisfied: python-dotenv>=0.13 in /usr/local/lib/python3.10/dist-packages (from uvicorn->onetwo==0.1.0) (1.0.1)\n",
            "Requirement already satisfied: uvloop!=0.15.0,!=0.15.1,>=0.14.0 in /usr/local/lib/python3.10/dist-packages (from uvicorn->onetwo==0.1.0) (0.19.0)\n",
            "Requirement already satisfied: watchfiles>=0.13 in /usr/local/lib/python3.10/dist-packages (from uvicorn->onetwo==0.1.0) (0.21.0)\n",
            "Requirement already satisfied: websockets>=10.4 in /usr/local/lib/python3.10/dist-packages (from uvicorn->onetwo==0.1.0) (12.0)\n",
            "Requirement already satisfied: httplib2<1dev,>=0.15.0 in /usr/local/lib/python3.10/dist-packages (from google-api-python-client->google-generativeai->onetwo==0.1.0) (0.22.0)\n",
            "Requirement already satisfied: google-auth-httplib2>=0.1.0 in /usr/local/lib/python3.10/dist-packages (from google-api-python-client->google-generativeai->onetwo==0.1.0) (0.1.1)\n",
            "Requirement already satisfied: uritemplate<5,>=3.0.1 in /usr/local/lib/python3.10/dist-packages (from google-api-python-client->google-generativeai->onetwo==0.1.0) (4.1.1)\n",
            "Requirement already satisfied: google-crc32c<2.0dev,>=1.0 in /usr/local/lib/python3.10/dist-packages (from google-resumable-media<3.0dev,>=0.6.0->google-cloud-bigquery!=3.20.0,<4.0.0dev,>=1.15.0->google-cloud-aiplatform->onetwo==0.1.0) (1.5.0)\n",
            "Requirement already satisfied: pyparsing!=3.0.0,!=3.0.1,!=3.0.2,!=3.0.3,<4,>=2.4.2 in /usr/local/lib/python3.10/dist-packages (from httplib2<1dev,>=0.15.0->google-api-python-client->google-generativeai->onetwo==0.1.0) (3.1.2)\n",
            "Requirement already satisfied: ml-dtypes>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from jax>=0.4.19->flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (0.2.0)\n",
            "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.10/dist-packages (from jax>=0.4.19->flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (3.3.0)\n",
            "Requirement already satisfied: scipy>=1.9 in /usr/local/lib/python3.10/dist-packages (from jax>=0.4.19->flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (1.11.4)\n",
            "Requirement already satisfied: pyasn1<0.7.0,>=0.4.6 in /usr/local/lib/python3.10/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3.0.0dev,>=2.14.1->google-cloud-aiplatform->onetwo==0.1.0) (0.6.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0.dev0,>=2.18.0->google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0dev,>=1.34.1->google-cloud-aiplatform->onetwo==0.1.0) (3.3.2)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0.dev0,>=2.18.0->google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0dev,>=1.34.1->google-cloud-aiplatform->onetwo==0.1.0) (2.0.7)\n",
            "Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.10/dist-packages (from rich>=11.1->flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (3.0.0)\n",
            "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.10/dist-packages (from rich>=11.1->flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (2.16.1)\n",
            "Requirement already satisfied: shellingham>=1.3.0 in /usr/local/lib/python3.10/dist-packages (from typer>=0.12.3->fastapi-cli>=0.0.2->fastapi->onetwo==0.1.0) (1.5.4)\n",
            "Requirement already satisfied: chex>=0.1.86 in /usr/local/lib/python3.10/dist-packages (from optax->flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (0.1.86)\n",
            "Requirement already satisfied: jaxlib>=0.1.37 in /usr/local/lib/python3.10/dist-packages (from optax->flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (0.4.26+cuda12.cudnn89)\n",
            "Requirement already satisfied: etils[epath,epy] in /usr/local/lib/python3.10/dist-packages (from orbax-checkpoint->flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (1.7.0)\n",
            "Requirement already satisfied: nest_asyncio in /usr/local/lib/python3.10/dist-packages (from orbax-checkpoint->flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (1.6.0)\n",
            "Requirement already satisfied: toolz>=0.9.0 in /usr/local/lib/python3.10/dist-packages (from chex>=0.1.86->optax->flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (0.12.1)\n",
            "Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.10/dist-packages (from markdown-it-py>=2.2.0->rich>=11.1->flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (0.1.2)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from etils[epath,epy]->orbax-checkpoint->flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (2023.6.0)\n",
            "Requirement already satisfied: importlib_resources in /usr/local/lib/python3.10/dist-packages (from etils[epath,epy]->orbax-checkpoint->flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (6.4.0)\n",
            "Requirement already satisfied: zipp in /usr/local/lib/python3.10/dist-packages (from etils[epath,epy]->orbax-checkpoint->flax>=0.8->gemma@ git+https://github.com/google-deepmind/gemma.git->onetwo==0.1.0) (3.18.2)\n"
          ]
        }
      ],
      "source": [
        "!pip install git+https://github.com/google-deepmind/onetwo"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "blfwJXpeHxPS"
      },
      "source": [
        "### Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "50CME0WvHy55"
      },
      "outputs": [],
      "source": [
        "from onetwo import ot\n",
        "from onetwo.builtins import llm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C-8S-I2fIBIn"
      },
      "source": [
        "### Register Gemma as the backend"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ig8q7ZcMHSjC"
      },
      "outputs": [],
      "source": [
        "from onetwo.backends import gemma_local\n",
        "\n",
        "backend = gemma_local.Gemma(checkpoint_path=checkpoint_path, vocab_path=vocab_path)\n",
        "backend.register()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I-hmUfo2IKyH"
      },
      "source": [
        "### Run inference with `llm.generate_text()`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PIAy8pXHIRlK"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            " the cities of <strong>Duisburg</strong>, <strong>Cologne</strong> and <strong>Bonn</strong>.\n",
            "\n",
            "<strong>Duisburg</strong> is a city\n"
          ]
        }
      ],
      "source": [
        "# Ask the model to continue the given text.\n",
        "e = llm.generate_text(\n",
        "    \"Three not so well known cities in Gemany are\", max_tokens=30, temperature=0.3\n",
        ")\n",
        "print(ot.run(e))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "[Gemma_1]Using_with_OneTwo.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
